
import numpy as np
from collections import defaultdict

class LandmarkSmoother:
    """
    指数平滑：用于平滑每个关节点坐标。
    alpha: 平滑系数 0<alpha<=1，值越小越平滑但延迟越大
    """

    def __init__(self, alpha: float = 0.7):
        self.alpha = alpha
        self.prev = defaultdict(list)  # person_id -> list of np.ndarray

    def smooth(self, persons):
        """
        persons: List[List[landmark]] - 每个人一组 landmarks
        return: 同样结构的平滑后 landmarks（复制）
        """
        smoothed = []
        for pid, lms in enumerate(persons):
            if not lms:
                smoothed.append(lms)
                continue

            prev_lms = self.prev.get(pid)
            curr = np.array([[lm.x, lm.y, lm.z, getattr(lm, "visibility", 1.0)] for lm in lms], dtype=np.float32)

            if prev_lms is None or len(prev_lms) != len(curr):
                self.prev[pid] = curr
                smoothed.append(lms)
                continue

            new_vals = self.alpha * curr + (1 - self.alpha) * prev_lms
            self.prev[pid] = new_vals

            # 还原成 MediaPipe 的 landmark-like 结构
            for i, lm in enumerate(lms):
                lm.x, lm.y, lm.z = new_vals[i, :3]
                if hasattr(lm, "visibility"):
                    lm.visibility = new_vals[i, 3]
            smoothed.append(lms)

        return smoothed


import numpy as np
from collections import defaultdict


class AdaptiveLandmarkSmoother:
    """
    自适应指数平滑 + 静止检测：
    - 动作慢：增加平滑
    - 动作快：减少平滑（提高响应）
    - 静止：冻结坐标（返回上帧结果）

    参数:
        alpha_min: 最低平滑系数
        alpha_max: 最高平滑系数
        sensitivity: 速度到alpha映射的灵敏度
        freeze_threshold: 判定为“静止”的平均速度阈值
        freeze_patience: 连续静止帧数，触发冻结
    """

    def __init__(
        self,
        alpha_min=0.4,
        alpha_max=0.9,
        sensitivity=3.0,
        freeze_threshold=0.002,
        freeze_patience=3,
    ):
        self.alpha_min = alpha_min
        self.alpha_max = alpha_max
        self.sensitivity = sensitivity
        self.freeze_threshold = freeze_threshold
        self.freeze_patience = freeze_patience

        # 状态缓存
        self.prev = defaultdict(lambda: None)
        self.freeze_counter = defaultdict(int)
        self.is_frozen = defaultdict(bool)

    def _compute_speed(self, prev, curr):
        if prev is None or len(prev) != len(curr):
            return 0.0
        diff = curr - prev
        return float(np.mean(np.linalg.norm(diff, axis=1)))

    def _adaptive_alpha(self, speed):
        speed = np.clip(speed, 0, 0.1)
        alpha = self.alpha_min + self.sensitivity * speed
        return float(np.clip(alpha, self.alpha_min, self.alpha_max))

    def _update_freeze_state(self, person_id, speed):
        """维护每个目标的冻结状态"""
        if speed < self.freeze_threshold:
            self.freeze_counter[person_id] += 1
            if self.freeze_counter[person_id] >= self.freeze_patience:
                self.is_frozen[person_id] = True
        else:
            # 一旦速度上升, 立即解除冻结
            self.freeze_counter[person_id] = 0
            self.is_frozen[person_id] = False

    def smooth(self, persons):
        smoothed = []
        for pid, lms in enumerate(persons):
            if not lms:
                smoothed.append(lms)
                continue

            curr = np.array([[lm.x, lm.y, lm.z] for lm in lms], dtype=np.float32)
            prev = self.prev[pid]
            speed = self._compute_speed(prev, curr)
            self._update_freeze_state(pid, speed)

            # 若被冻结，直接返回上一帧
            if self.is_frozen[pid] and prev is not None:
                frozen_vals = prev.copy()
                for i, lm in enumerate(lms):
                    lm.x, lm.y, lm.z = frozen_vals[i, :3]
                smoothed.append(lms)
                continue

            alpha = self._adaptive_alpha(speed)
            if prev is None:
                self.prev[pid] = curr
                smoothed.append(lms)
                continue

            new_vals = alpha * curr + (1 - alpha) * prev
            self.prev[pid] = new_vals
            for i, lm in enumerate(lms):
                lm.x, lm.y, lm.z = new_vals[i, :3]
            smoothed.append(lms)

        return smoothed
